import math
import matplotlib.pyplot as plt
import numpy as np

plt.rcParams['font.sans-serif'] = ['SimHei']     
plt.rcParams['axes.unicode_minus'] = False        


#距离函数
def euclidean_distance(x1, x2):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(x1, x2)))


#KNN 预测函数 
def knn_predict(train_X, train_y, x, k=3):
    distances = []
    for xi, yi in zip(train_X, train_y):
        d = euclidean_distance(x, xi)
        distances.append((d, yi, xi))
    
    distances.sort(key=lambda t: t[0])
    k_nearest = distances[:k]

    labels = [label for _, label, _ in k_nearest]
    pred = max(set(labels), key=labels.count)
    return pred, k_nearest


# 自动颜色字典
def color_map(labels):
    unique = sorted(set(labels))
    colors = plt.cm.tab10(np.linspace(0, 1, len(unique)))
    return {label: colors[i] for i, label in enumerate(unique)}


# 自动生成两类训练点
np.random.seed(2)

cluster_0 = np.random.randn(50, 2) * 0.6 + np.array([2, 2])
cluster_1 = np.random.randn(50, 2) * 0.6 + np.array([6, 7])

train_X = np.vstack([cluster_0, cluster_1])
train_y = [0]*50 + [1]*50

# 多个测试点
test_points = [
    [3, 3],
    [6, 5.5],
    [4, 6],
]
k = 5

# 可视化
plt.figure(figsize=(9, 9))

colors = color_map(train_y)

for (x, y), label in zip(train_X, train_y):
    plt.scatter(x, y, s=60, color=colors[label], edgecolor='black')

for label, color in colors.items():
    plt.scatter([], [], color=color, label=f"Class {label}")
plt.legend(fontsize=12)


for idx, test_point in enumerate(test_points):
    pred_label, nearest = knn_predict(train_X, train_y, test_point, k)

    # 测试点
    plt.scatter(test_point[0], test_point[1],
                color='none', edgecolor='green',
                s=260, linewidths=2.5, marker='^')
    plt.text(test_point[0] + 0.08, test_point[1] + 0.08,
             f"Test {idx+1} (Pred={pred_label})",
             fontsize=12, color="green")

    # 画 K 临近点 + 距离
    for i, (dist, label, p) in enumerate(nearest):
        plt.plot([test_point[0], p[0]], [test_point[1], p[1]],
                 linestyle="--", linewidth=1.2)

        mid_x = (test_point[0] + p[0]) / 2
        mid_y = (test_point[1] + p[1]) / 2
        plt.text(mid_x, mid_y, f"{dist:.2f}", fontsize=10)

        plt.scatter(p[0], p[1], s=180, facecolors='none',
                    edgecolors='black', linewidths=2)
        plt.text(p[0] - 0.1, p[1] - 0.15,
                 f"#{i+1}", fontsize=12, color='black')


plt.gca().set_aspect('equal', adjustable='box')
plt.title(f"KNN 可视化（k={k}） - 自动生成两类聚类点")
plt.xlabel("X1")
plt.ylabel("X2")
plt.grid(True)
plt.show()
